package main
import (
	"fmt"
	"net"
	"os"
	"errors"
	"strings"
)
const  (
	QueryType_ARECORD=1
	QueryType_NS=2
	QueryType_CNAME=5
	QueryType_PTR=12
	QueryType_HINFO=13
	QueryType_MX=15
	QueryType_TXT=16
	QueryType_AXFR=252
	QueryType_ANY=255
	
	QueryClass_RESERVED=0
	QueryClass_INTERNET=1
	QueryClass_UNASSINGED=2
	QueryClass_CHAOS=3
	QueryClass_HESIOD=4
)

type Flags struct {
    QR  uint8// 1 查询/响应的标志位，1为响应，0为查询  
    Opcode  uint8// 4 定义查询或响应的类型（0为标准的，1为反向的，2为服务器状态请求，3-15保留值） 
    AA uint8// 1  授权回答的标志位。该位在响应报文中有效，1表示域名服务器是权限服务器 
    TC uint8// 1  截断标志位。1表示响应已超过512字节并已被截断 
    RD uint8//1  被请求报文设置，该位为1表示客户端希望得到递归回答，应答时使用相同的值返回。 
    RA uint8//1 支持递归， 这个比特位在应答中设置或取消，用来代表服务器是否支持递归查询。 
    Zero uint8//3 保留字段。 
    Rcode uint8//4 返回码，表示响应的差错状态，通常为0和3，各取值含义如下： 
    /*
    0 无差错 
	1 报文格式差错(Format error)服务器不能理解的请求报文 
	2 服务器失败（Server failure）因为服务器的原因导致没办法处理的请求 
	3 名字错误（Name error） 只有对授权域名解析服务器有意义，指出解析的域名不存在 
	4 没有实现（Not Implemented）域名服务器不支持的查询类型 
	5 拒绝（Refused）服务器由于设置的策略据局给出应答 
	6 - 15 保留值 
    */
}

func (f *Flags)Encode()[]byte{
	s := make([]byte,2)
	var a uint8 = 0
	a=(a|f.QR)<<7
	a=a|((a|f.Opcode)<<3)
	a=a|((a|f.AA)<<1)
	a=(a|f.TC)
	s[0]=byte(a)
	a=0
	a=(a|f.RD)<<7
	a=a|((a|f.RA)<<6)
	a=a|((a|f.Zero)<<3)
	a=(a|f.Rcode)
	s[1]=byte(a)
	return s
}

type PackageHeader struct{
	Id uint16
	Flags *Flags
	Questions uint16
	Answers uint16
	Nameservers uint16
	AdditionalRecords uint16
}
func (f *PackageHeader)Encode()[]byte{
	s := make([]byte,12)
	s[0]=byte(f.Id>>8)
	s[1]=byte((f.Id<<8)>>8)
	tmp := f.Flags.Encode()
	s[2]=tmp[0]
	s[3]=tmp[1]
	s[4]=byte(f.Questions>>8)
	s[5]=byte((f.Questions<<8)>>8)
	s[6]=byte(f.Answers>>8)
	s[7]=byte((f.Answers<<8)>>8)
	s[8]=byte(f.Nameservers>>8)
	s[9]=byte((f.Nameservers<<8)>>8)
	s[10]=byte(f.AdditionalRecords>>8)
	s[11]=byte((f.AdditionalRecords<<8)>>8)
	return s
}
type Record struct{
	Name string
    Qtype uint16
    Qclass uint16
}
func (f *Record)Encode()[]byte{
	l:=computeRecordLen(f)
	s := make([]byte,l)
	splits := strings.Split(strings.Trim(f.Name,"."),".")
	idx := 0
	for i:=0;i<len(splits);i+=1{
		s[idx] = byte(len(splits[i]))
		idx+=1
		for n:=0;n<len(splits[i]);n+=1{
			s[idx] = splits[i][n]
			idx+=1
		}
	}
	s[idx]=byte(0)
	s[idx+1]=byte(f.Qtype>>8)
	s[idx+2]=byte((f.Qtype<<8)>>8)
	s[idx+3]=byte(f.Qclass>>8)
	s[idx+4]=byte((f.Qclass<<8)>>8)
	return s
}

type RRRecord struct{
	Name string
    Qtype uint16
	Qclass uint16
	TTL   uint32
	DataLen  uint16
	Data []byte
}

func (f *RRRecord)Encode()[]byte{
	l:=computeRRRecordLen(f)
	s := make([]byte,l)
	splits := strings.Split(strings.Trim(f.Name,"."),".")
	idx := 0
	for i:=0;i<len(splits);i+=1{
		s[idx] = byte(len(splits[i]))
		idx+=1
		for n:=0;n<len(splits[i]);n+=1{
			s[idx] = splits[i][n]
			idx+=1
		}
	}
	s[idx]=byte(0)
	s[idx+1]=byte(f.Qtype>>8)
	s[idx+2]=byte((f.Qtype<<8)>>8)
	s[idx+3]=byte(f.Qclass>>8)
	s[idx+4]=byte((f.Qclass<<8)>>8)

	s[idx+5]=byte(f.TTL>>8)
	s[idx+6]=byte((f.TTL<<8)>>8)
	s[idx+7]=byte(f.TTL>>8)
	s[idx+8]=byte((f.TTL<<8)>>8)
	s[idx+9]=byte(f.DataLen>>8)
	s[idx+10]=byte((f.DataLen<<8)>>8)

	idx+=11
	for i:= 0;i<int(f.DataLen);i+=1{
		s[idx+i] = f.Data[i]
	}
	return s
}

type Query struct {
	Header *PackageHeader 
    Questions []*Record
    Answers []*RRRecord
    Nameservers []*RRRecord
    AdditionalRecords []*RRRecord
}

func (f *Query)Encode()[]byte{
	l:=computePackageLen(f)
	s := make([]byte,l)
	he:= f.Header.Encode()
	idx := 0
	for i:=0;i<len(he);i+=1{
		s[idx] = he[i]
		idx+=1
	}
	for n:=0;n<int(f.Header.Questions);n+=1{
		he = f.Questions[n].Encode()
		for i:=0;i<len(he);i+=1{
			s[idx] = he[i]
			idx+=1
		}
	}
	for n:=0;n<int(f.Header.Answers);n+=1{
		he = f.Answers[n].Encode()
		for i:=0;i<len(he);i+=1{
			s[idx] = he[i]
			idx+=1
		}
	}
	for n:=0;n<int(f.Header.Nameservers);n+=1{
		he = f.Nameservers[n].Encode()
		for i:=0;i<len(he);i+=1{
			s[idx] = he[i]
			idx+=1
		}
	}
	for n:=0;n<int(f.Header.AdditionalRecords);n+=1{
		he = f.AdditionalRecords[n].Encode()
		for i:=0;i<len(he);i+=1{
			s[idx] = he[i]
			idx+=1
		}
	}
	return s
}


func ParseFlags(query []byte) (*Flags,error) {
	if len(query) < 2 {
		return nil,errors.New("data is too short.")
	}
	ret := &Flags{}
	ret.QR=uint8(query[0])>>7
	ret.Opcode=(uint8(query[0])&120)>>3
	ret.AA=(uint8(query[0])&7)>>2
	ret.TC=(uint8(query[0])&3)>>1
	ret.RD=(uint8(query[0])&1)
	ret.RA=uint8(query[1])>>7
	ret.Zero=(uint8(query[1])&112)>>4
	ret.Rcode=(uint8(query[1])&15)
	return ret,nil
}

func ParsePackageHeader(query []byte) (*PackageHeader,error) {
	if len(query) < 12 {
		return nil,errors.New("data is too short.")
	}
	ret := &PackageHeader{}
	ret.Id = uint16(query[0])<<8+uint16(query[1])
	var err error=nil
	ret.Flags, err = ParseFlags(query[2:4])
	if(err !=nil){
		return nil ,err
	}
	ret.Questions = uint16(query[4])<<8+uint16(query[5])
	ret.Answers = uint16(query[6])<<8+uint16(query[7])
	ret.Nameservers = uint16(query[8])<<8+uint16(query[9])
	ret.AdditionalRecords = uint16(query[10])<<8+uint16(query[11])
	
	return ret,nil

}

func ParseRecord(query []byte) (*Record,uint16,error) {
	r := &Record{}
	var pos uint16=0
	name := ""
	for{
		l:=uint16(query[pos])
		if l == 0 {
			break
		}
		if len(query) < int(pos+1+l){
			return nil,0,errors.New("data is too short.")
		}
		name +=string(query[pos+1:pos+1+l])
		name +="."
		pos+=l+1
	}
	pos+=1
	if len(query) < int(pos+4){
		return nil,0,errors.New("data is too short.")
	}
	r.Name = name
	r.Qtype = uint16(query[pos])<<8+uint16(query[pos+1])
	r.Qclass = uint16(query[pos+2])<<8+uint16(query[pos+3])

	return r,pos+4,nil
}

func ParseRecords(query []byte, num uint16) ([]*Record,error) {
	r := make([]*Record,0)
	var pos uint16=0
	var item *Record=nil
	var err error =nil
	for{
		var i uint16=0
		if num<=0 {
			break
		}
		item,i,err = ParseRecord(query[pos:])
		pos += i
		//fmt.Printf("record:%+v,pos:%v,err:%+v\n",item,pos,err)
		if err !=nil {
			return nil,err
		}
		r=append(r,item)
		num-=1
	}
	return r,nil
}

func ParseQuery(query []byte) (*Query,error) {
	q := &Query{}
	var err error=nil
	q.Header,err  = ParsePackageHeader(query)
	//fmt.Printf("header:%+v,err:%+v\n",q.Header,err)
	if(err !=nil){
		return nil,err
	}
	q.Questions,err = ParseRecords(query[12:],q.Header.Questions)
	return q,nil
}



func computeRecordLen(rcd *Record)uint16{
	var r uint16=4
	r+=2+uint16(len(rcd.Name))
	return r
}
func computeRRRecordLen(rcd *RRRecord)uint16{
	var r uint16=10
	r+=2+uint16(len(rcd.Name))
	r+=rcd.DataLen+2
	return r
}
func computePackageLen(pack *Query)uint16{
	var r uint16=12
	if pack.Header.Questions >0 {
		for i :=uint16(0);i<pack.Header.Questions;i+=1{
			r+=computeRecordLen(pack.Questions[i])
		}
	}
	if pack.Header.Answers >0 {
		for i :=uint16(0);i<pack.Header.Answers;i+=1{
			r+=computeRRRecordLen(pack.Answers[i])
		}
	}
	if pack.Header.Nameservers >0 {
		for i :=uint16(0);i<pack.Header.Nameservers;i+=1{
			r+=computeRRRecordLen(pack.Nameservers[i])
		}
	}
	if pack.Header.AdditionalRecords >0 {
		for i :=uint16(0);i<pack.Header.AdditionalRecords;i+=1{
			r+=computeRRRecordLen(pack.AdditionalRecords[i])
		}
	}
	return r
}

func response(request *Query)(*Query,error){
	r := request
	r.Header.Flags.QR = 1
	r.Header.Answers=1
	answer:=&RRRecord{}
	answer.Name=r.Questions[0].Name
	answer.TTL=1024
	answer.Qclass=r.Questions[0].Qclass
	//answer.Qtype = r.Questions[0].Qtype
	answer.Qtype = 1
	answer.DataLen=8
	answer.Data=[]byte{127,0,0,1,192,168,0,1}
	r.Answers = append(r.Answers,answer)
	return r,nil
}

func handle(conn *net.UDPConn,query []byte,client *net.UDPAddr){
	
	q,e := ParseQuery(query)
	if(e !=nil){
		fmt.Println("error: ",e)
		return
	}
	//fmt.Printf("query: %#v\n",q)
	r,err := response(q)
	if(err !=nil){
		fmt.Println("error: ",e)
		return
	}

	conn.WriteToUDP(r.Encode(), client)
}


func main(){
	address := "127.0.0.1:53"
	addr, err := net.ResolveUDPAddr("udp", address)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
 
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
 
	defer conn.Close()
 
	for {
		// Here must use make and give the lenth of buffer
		data := make([]byte, 1024)
		l, rAddr, err := conn.ReadFromUDP(data)
		fmt.Println(l,rAddr,err)
		if err != nil {
			fmt.Println(err)
			continue
		}
		go handle(conn, data[0:l] , rAddr)
 
	}

}
